'''
@Author  : Pinde Song    
@Email   : tsingachieve@163.com
@Time    : 2022/10/10 20:06
@File    : visdrone2coco.py
'''
import os
from tqdm import tqdm
from PIL import Image
import json

def visdrone2coco(root_path, type):
    # 初始化节点信息
    annotations = []
    annotation_id = 1
    # 初始化图片信息
    images = []
    image_id = 1
    # 全局字典
    global_dict = {}
    # 初始化类别信息
    categories = [
        {"id": 1, "name": "pedestrian"}, {"id": 2, "name": "people"}, {"id": 3, "name": "bicycle"}, {"id": 4, "name": "car"},
        {"id": 5, "name": "van"},{"id": 6, "name": "truck"}, {"id": 7, "name": "tricycle"}, {"id": 8, "name": "awning-tricycle"},
        {"id": 9, "name": "bus"},{"id": 10, "name": "motor"}
    ]
    # 读取txt文件
    annotation_path = os.path.join(root_path, "annotations")
    ans = os.listdir(annotation_path)
    # 初始化进度条
    ans_bar = tqdm(ans)
    for index, file_name in enumerate(ans_bar):
        # 获取图片数据信息
        image_path = os.path.join(root_path, "images")
        image_source = Image.open(os.path.join(image_path, file_name.strip(".txt") + ".jpg"), "r")
        width, height = image_source.size
        image_dict = {
            "file_name": file_name.strip(".txt") + ".jpg",
            "width": width,
            "height": height,
            "id": image_id
        }

        loc = os.path.join(annotation_path, file_name)
        with open(loc, "r") as f:
            # top_x, top_y, w, h, score, category, ...
            for file_content in f.readlines():
                label_msg = file_content.strip("\n").split(",")
                if label_msg[5] not in ["0", "11"]:
                    label_dict = {
                        "id": annotation_id,
                        "image_id": image_id,
                        "category_id": int(label_msg[5]),
                        "bbox": list(map(lambda x: int(x), label_msg[:4])),
                        "area": int(label_msg[2]) * int(label_msg[3]),
                        "iscrowd": 0
                    }
                    # 标签信息存入节点
                    annotations.append(label_dict)
                    annotation_id += 1
        f.close()

        # 图片信息存入节点
        images.append(image_dict)
        # 更新图片id
        image_id += 1
        # 设置进度信息
        ans_bar.set_description("[{}/{}]converting ...".format(index, len(ans)))

    # 将所有数据保存至全局
    global_dict["images"] = images
    global_dict["annotations"] = annotations
    global_dict["categories"] = categories
    # 保存json文件
    with open("./{}-instances.json".format(type), "w") as f:
        json.dump(global_dict, f)
    f.close()

visdrone2coco("./datasets/VisDrone2019-DET-train/", "train")
visdrone2coco("./datasets/VisDrone2019-DET-val/", "val")